from pylab import *

class SearchablePoints:
    def __init__(o, width, height, r0):
        o.width = width
        o.height = height
        o.r0 = r0
        o.xbins = int(o.width / o.r0) + 1
        o.ybins = int(o.height / o.r0) + 1
        # print "bin count: %d x %d"%(o.xbins, o.ybins)

        o.boxes = [[ [] for y in range(o.ybins) ]
                    for x in range(o.xbins) ]
        o.points = []

    def add_point(o, point):
        x, y = point
        nx = int(x/o.r0)
        ny = int(y/o.r0)
        o.points.append(point)
        o.boxes[nx][ny].append(point)

    def find_nearest(o, point):
        """finds all points that are at most 1 box distance away from point."""
        result = []
        x, y = point
        nx = int(x/o.r0)
        ny = int(y/o.r0)
        nnx = [nx-1, nx, nx+1]
        nny = [ny-1, ny, ny+1]
        if nx==0: nnx.remove(nx-1)
        if ny==0: nny.remove(ny-1)
        if nx==o.xbins-1: nnx.remove(nx+1)
        if ny==o.ybins-1: nny.remove(ny+1)

        for nx in nnx:
            for ny in nny:
                for pt in o.boxes[nx][ny]:
                    px, py = pt
                    if hypot(x-px, y-py) <= r0:
                        result.append(pt)
        try:
            result.remove(point)
        except:
            pass
        return result


def showradius(event):
    print "point: ", event.xdata, event.ydata
    arg = linspace(0, 2*pi, 50)
    cx = r0 * cos(arg) + event.xdata
    cy = r0 * sin(arg) + event.ydata
    plot(cx, cy, "g-")
    points = sp.find_nearest((event.xdata, event.ydata))
    points = array(points)
    plot(points[:,0], points[:,1], "rx")

def do_step(spoints, factor = 1.0):
    width, height, r0 = spoints.width, spoints.height, spoints.r0
    result = SearchablePoints(width, height, r0)
    for point in spoints.points:
        others = spoints.find_nearest(point)
        forcex, forcey = 0., 0.
        for other in others:
            dist = hypot(other[0]-point[0], other[1]-point[1])
            if dist > r0: continue
            forcex += (point[0] - other[0]) / dist * (1 - dist/r0)
            forcey += (point[1] - other[1]) / dist * (1 - dist/r0)
        forcex *= 0.5 * r0 * factor
        forcey *= 0.5 * r0 * factor
        x = point[0] + forcex
        y = point[1] + forcey
        if x<0: x=0.0
        if y<0: y=0.0
        if x>width: x=width;
        if y>height: y=height
        result.add_point((x, y))
    return result
            
count= 25
r0 = 1.5*sqrt(400*300/count)
sp = SearchablePoints(400, 300, r0)
step = 0
traj_x = 0
traj_y = 0
show_trajectory_points = 10

def main():
    global traj_x, traj_y
    points = rand(count, 2)
    points[:,0] *= 400
    points[:,1] *= 300
    [sp.add_point(tuple(pt)) for pt in points]
    points = array(sp.points)
    traj_x = points[:,0].reshape(1, count)
    traj_y = points[:,1].reshape(1, count)

    data = plot(traj_x, traj_y, "r-o")
    n = 0
    def next_step(event):
        global sp 
        global step
        global traj_x, traj_y
        if step >= 19:
            factor=0.5
        else:
            factor = 1.0/(20-step)
        sp = do_step(sp, factor)
        points = array(sp.points)
        traj_x = concatenate((traj_x, points[:, 0].reshape(1, count)), axis=0)
        traj_y = concatenate((traj_y, points[:, 1].reshape(1, count)), axis=0)
        for n in xrange(len(data)):
            data[n].set_data(traj_x[-show_trajectory_points:, n], traj_y[-show_trajectory_points:,n])
        draw()
        step += 1
        print "step: %d"%step
    connect("button_press_event", next_step)
    gca().set_aspect("equal")
    show()


if __name__=="__main__": main()
